from config import GetConfig
import torch
from datalist import DataSetRace
from torch.utils.data import DataLoader
from model import Models,mobilenet_v3_small,resnet18
from torch import nn
from tqdm import tqdm
gpu_id = 0
class Train:
    def __init__(self) -> None:
        self.args = GetConfig()
        print(f'-----------{self.args.project_name}-----------')
        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        self.device = torch.device(f"cuda:{gpu_id}" if use_cuda else "cpu")

        kwargs = {"num_workers": 2, "pin_memory": True} if use_cuda else {
            "num_workers": 2, "pin_memory": True}
        
        datasetrace = DataSetRace(
            self.args.data_dir
        )
        datasettest = DataSetRace(
            r'D:\CartoonCameraProject\test_data\black_person\black_person',
            default_label=1
        )
        self.train_dataloader = DataLoader(
            datasetrace,
            batch_size=self.args.train_batch_size,
            shuffle=True,
            **kwargs,
        )
        self.test_dataloader = DataLoader(
            datasettest,
            batch_size=50,
            shuffle=False,
            **kwargs,
        )
        
        
        self.model = Models(num_class=2)
        # self.model = mobilenet_v3_small(num_classes=2)

        if use_cuda:
            self.model = torch.nn.DataParallel(self.model,device_ids=[gpu_id])
            
        if self.args.resume:
            self.model.load_state_dict(torch.load(self.args.pretrained_weight,map_location=self.device)['model_state_dict'],False)
            
            
        self.optim = torch.optim.SGD(
            self.model.parameters(),
            lr=self.args.lr,
            weight_decay=1e-4,
        )
        
        
        self.loss = nn.CrossEntropyLoss()
    def work(self):
        for epoch in range(self.args.epochs):
            self.train(epoch)
            self.test(epoch)
            torch.save({
                'model_state_dict':self.model.state_dict()
            },
                self.args.pretrained_weight)
    def train(self,epoch):
        self.model.train()
        pbar = tqdm(self.train_dataloader,desc=f'Train Epoch:{epoch}/{self.args.epochs}')
        correct = torch.zeros(1).squeeze().to(self.device)
        total = torch.zeros(1).squeeze().to(self.device)
        for data ,gt in pbar:
            data = data.to(self.device)
            gt = gt.to(self.device)
            
            self.optim.zero_grad()
            output = self.model(data)
            loss = self.loss(output,gt)
            loss.backward()
            
            self.optim.step()
            
            pred = torch.argmax(output,1)
            correct += (pred == gt).sum().float()
            total += len(gt)
            predict_acc = correct / total
            
            pbar.set_description(
                '"\033[0;31;40m"'+-+
                
                .-*  
                .           
                .
                f'Loss:{loss.item()} ||'
                f'ACC:{predict_acc}'
                + "\033[0m"
            )
    @torch.no_grad()
    def test(self,epoch):
        self.model.eval()
        pbar = tqdm(self.test_dataloader,desc=f'Test Epoch: {epoch}/{self.args.epochs}')
        correct = torch.zeros(1).squeeze().to(self.device)
        total = torch.zeros(1).squeeze().to(self.device)
        for data ,gt in pbar:
            data = data.to(self.device)
            gt = gt.to(self.device)
            output = self.model(data)
            pred = torch.argmax(output,1)
            correct += (pred == gt).sum().float()
            total += len(gt)
            predict_acc = correct / total
            pbar.set_description(
                "\033[1;32m"+
                f'Train Epoch:{epoch}/{self.args.epochs} ||'
                f'ACC:{predict_acc}'
                + "\033[0m"
            )
if __name__ == "__main__":
       model = Train()
       model.work()